import numpy as np
import os

os.environ["KERAS_BACKEND"] = "torch"  # 使用torch 作为后端
from keras.models import Sequential, load_model
from keras.layers import Dense, Dropout, Input
from keras.losses import MeanSquaredError
from sklearn.model_selection import train_test_split
from keras.callbacks import ModelCheckpoint
from utils.load_data import Pose
from sklearn.metrics import r2_score

# 加载数据
pose_data_list = Pose.load_data("C:/Users/kang_/Desktop/data")
X = np.array(
    [i.body_size(512 / i.img_width) for i in pose_data_list if i.weight > 0],
    dtype=np.float32,
)
Y = np.array([i.weight for i in pose_data_list if i.weight > 0], dtype=np.float32)
X_train, X_val, Y_train, Y_val = train_test_split(X, Y, test_size=0.1)

# 搭建模型
model = Sequential()
model.add(Input(shape=(X[0].size,)))
model.add(Dense(512, activation="relu"))  # 输入层到隐藏层
model.add(Dense(1024, activation="relu"))  # 隐藏层
model.add(Dropout(0.2))
model.add(Dense(1))  # 输出层
# 编译模型
model.compile(optimizer="adam", loss=MeanSquaredError(), metrics=["accuracy"])
model.summary()
checkpoint = ModelCheckpoint(
    "best_model.keras", monitor="val_loss", save_best_only=True, mode="min"
)

history = model.fit(
    X_train,
    Y_train,
    epochs=500,
    batch_size=20,
    verbose=2,
    validation_split=0.2,
    callbacks=[checkpoint],
)
model_best = load_model("./best_model_v1.keras")
pred = model_best.predict(X_val)
print(r2_score(Y_val, pred))
